import cv2
import albumentations as A
from albumentations.pytorch import ToTensorV2
import matplotlib.pyplot as plt
import os
import numpy as np

def get_train_transform():
    """训练时的数据增强transforms"""
    return A.Compose([
        # 几何变换
        A.Resize(height=128, width=128, p=1.0),  # 先resize到稍大尺寸
        A.RandomCrop(height=112, width=112, p=1.0),  # 随机裁剪到目标尺寸
        A.HorizontalFlip(p=0.5),
        A.Rotate(limit=15, p=0.3),
        A.ShiftScaleRotate(
            shift_limit=0.1, 
            scale_limit=0.1, 
            rotate_limit=10, 
            p=0.3
        ),
        
        # 光照变换
        A.RandomBrightnessContrast(
            brightness_limit=0.2, 
            contrast_limit=0.2, 
            p=0.4
        ),
        A.RandomGamma(gamma_limit=(80, 120), p=0.3),
        A.CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=0.3),
        
        # 噪声和模糊
        A.GaussNoise(var_limit=(10.0, 50.0), p=0.2),
        A.GaussianBlur(blur_limit=(3, 7), p=0.2),
        A.MotionBlur(blur_limit=7, p=0.1),
        
        # 色彩变换
        A.HueSaturationValue(
            hue_shift_limit=10, 
            sat_shift_limit=15, 
            val_shift_limit=10, 
            p=0.3
        ),
        A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.2),
        A.ChannelShuffle(p=0.1),
        
        # 遮挡和擦除
        A.CoarseDropout(
            max_holes=8, 
            max_height=8, 
            max_width=8, 
            min_holes=1, 
            min_height=4, 
            min_width=4, 
            p=0.2
        ),
        A.GridDropout(ratio=0.2, p=0.1),
        
        # 最终处理
        # A.ToGray(p=1.0),  # 转换为灰度图
        A.Normalize(mean=[0.0], std=[1.0], max_pixel_value=255.0),  # 简单的0-1归一化，相当于/255.0
        ToTensorV2()
    ])


def get_val_transform():
    """验证时的基础transforms"""
    return A.Compose([
        A.Resize(height=112, width=112, p=1.0),
        # A.ToGray(p=1.0),
        A.Normalize(mean=[0.0], std=[1.0], max_pixel_value=255.0),  # 简单的0-1归一化，相当于/255.0
        ToTensorV2()
    ])


def get_visualization_transform():
    """用于可视化的transform（不包含ToTensor）"""
    return A.Compose([
        # 几何变换
        A.Resize(height=128, width=128, p=1.0),
        A.RandomCrop(height=112, width=112, p=1.0),
        A.HorizontalFlip(p=0.5),
        A.Rotate(limit=15, p=0.5),
        A.ShiftScaleRotate(
            shift_limit=0.1, 
            scale_limit=0.1, 
            rotate_limit=10, 
            p=0.5
        ),
        
        # 光照变换
        A.RandomBrightnessContrast(
            brightness_limit=0.2, 
            contrast_limit=0.2, 
            p=0.5
        ),
        A.RandomGamma(gamma_limit=(80, 120), p=0.3),
        A.CLAHE(clip_limit=2.0, tile_grid_size=(8, 8), p=0.3),
        
        # 噪声和模糊
        A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),
        A.GaussianBlur(blur_limit=(3, 7), p=0.3),
        
        # 色彩变换
        A.HueSaturationValue(
            hue_shift_limit=10, 
            sat_shift_limit=15, 
            val_shift_limit=10, 
            p=0.4
        ),
        
        # 遮挡
        A.CoarseDropout(
            max_holes=5, 
            max_height=8, 
            max_width=8, 
            min_holes=1, 
            min_height=4, 
            min_width=4, 
            p=0.3
        ),
        
        # 转换为灰度但不normalize（用于可视化）
        # A.ToGray(p=1.0),
    ])


def create_augmentation_grid(image_path, num_samples=100):
    """
    创建10x10的增强图像网格用于可视化
    
    Args:
        image_path: 输入图像路径
        num_samples: 生成样本数量（默认100个用于10x10网格）
    """
    # 读取原始图像
    original_image = cv2.imread(image_path)
    if original_image is None:
        print(f"无法读取图像: {image_path}")
        return
    
    # original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
    
    # 获取可视化transform
    transform = get_visualization_transform()
    
    # 生成增强图像
    augmented_images = []
    
    for i in range(num_samples):
        # 应用transform
        augmented = transform(image=original_image)
        aug_image = augmented['image']
        
        # 确保图像是uint8格式用于显示
        if aug_image.dtype != np.uint8:
            aug_image = (aug_image * 255).astype(np.uint8)
        
        augmented_images.append(aug_image)
    
    # 创建10x10网格
    grid_size = 10
    img_size = 112
    
    # 创建空白画布
    grid_image = np.zeros((grid_size * img_size, grid_size * img_size, 3), dtype=np.uint8)
    
    # 填充网格
    for i in range(grid_size):
        for j in range(grid_size):
            idx = i * grid_size + j
            if idx < len(augmented_images):
                # 如果是3通道图像，转换为灰度
                # if len(augmented_images[idx].shape) == 3:
                #     img = cv2.cvtColor(augmented_images[idx], cv2.COLOR_RGB2GRAY)
                # else:
                img = augmented_images[idx]
                
                # 放置到网格中
                y_start = i * img_size
                y_end = (i + 1) * img_size
                x_start = j * img_size
                x_end = (j + 1) * img_size
                
                grid_image[y_start:y_end, x_start:x_end] = img
    
    return grid_image


def test_transforms_performance():
    """测试不同transforms的性能"""
    print("=== Transform性能测试 ===")
    
    # 创建测试图像
    test_image = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
    
    transforms_to_test = {
        'train': get_train_transform(),
        'val': get_val_transform(),
        'viz': get_visualization_transform()
    }
    
    import time
    
    for name, transform in transforms_to_test.items():
        start_time = time.time()
        
        # 测试100次变换
        for _ in range(100):
            if name == 'viz':
                _ = transform(image=test_image)
            else:
                _ = transform(image=test_image)
        
        end_time = time.time()
        avg_time = (end_time - start_time) / 100 * 1000  # 转换为毫秒
        
        print(f"{name.upper()} Transform: {avg_time:.2f}ms per image")


def test_normalization_output():
    """测试归一化输出范围"""
    print("=== 测试归一化输出范围 ===")
    
    # 创建测试图像
    test_image = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
    
    # 测试训练transform
    train_transform = get_train_transform()
    val_transform = get_val_transform()
    
    # 应用transform
    train_result = train_transform(image=test_image)['image']
    val_result = val_transform(image=test_image)['image']
    
    print(f"训练Transform输出:")
    print(f"  Shape: {train_result.shape}")
    print(f"  Data type: {train_result.dtype}")
    print(f"  Min value: {train_result.min():.4f}")
    print(f"  Max value: {train_result.max():.4f}")
    print(f"  Mean value: {train_result.mean():.4f}")
    
    print(f"\n验证Transform输出:")
    print(f"  Shape: {val_result.shape}")
    print(f"  Data type: {val_result.dtype}")
    print(f"  Min value: {val_result.min():.4f}")
    print(f"  Max value: {val_result.max():.4f}")
    print(f"  Mean value: {val_result.mean():.4f}")


if __name__ == "__main__":
    # 测试数据集加载
    print("=== 测试数据集加载 ===")
    
    # 如果有真实图像路径，使用真实图像
    sample_image_path = "/Users/tunm/Downloads/AffectNet/validnew/113.jpg"  # 请替换为实际图像路径
    
    # 检查文件是否存在
    if os.path.exists(sample_image_path):
        print(f"使用真实图像: {sample_image_path}")
        grid_image = create_augmentation_grid(sample_image_path)
        
        if grid_image is not None:
            # 显示网格图像
            cv2.imshow('Augmentation Grid (10x10)', grid_image)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
            
            # 保存网格图像
            cv2.imwrite('augmentation_grid.png', grid_image)
            print("增强网格图像已保存为 'augmentation_grid.png'")
    else:
        print(f"图像文件不存在: {sample_image_path}")
        print("请检查路径是否正确")

    # 归一化测试
    print("\n=== 归一化测试 ===")
    test_normalization_output()

    # 性能测试
    print("\n=== 性能测试 ===")
    test_transforms_performance()
    
    print("\n=== 测试完成 ===")
    print("建议使用方式:")
    print("train_dataset = AffectDataSet2(data_path, 'train', 7, get_train_transform())")
    print("val_dataset = AffectDataSet2(data_path, 'val', 7, get_val_transform())")
    print("\n注意: 现在tensor数据范围是0-1，相当于pixel_value/255.0")